"""
OCR 文案对比可视化工具 (Tkinter)
功能：
- 选择图片文件夹和 CSV 文件夹（按文件名一一对应，例如 价格页.jpg <-> 价格页.csv）
- 对每张图片做 OCR（支持 pytesseract 或 easyocr）
- 将 OCR 识别到的文本与对应 CSV 中的每条文案逐条比对
  - 支持精确包含检测与模糊匹配（可调阈值）
- 实时日志输出（界面与文件）
- 显示每条文案的比对结果（通过/未通过、相似度分数）
- 支持导出比对结果为 CSV

改进：
- 增加对“白字在彩色按钮上”和“浅灰细字”两类常见难识别场景的区域增强预处理
- 自动尝试检测红色/彩色按钮区域并对按钮区域做反色+二值化识别
- 增加按钮形状候选检测（基于边缘/轮廓），以捕捉非标准红色按钮（如渐变、粉色）
- 对页面底部提示文本（默认取底部 25% 区域）做 CLAHE 增强+锐化再识别
- 对按钮区域尝试多种预处理和多种 OCR 参数组合，取并集以提高召回率

依赖：
- Pillow
- pytesseract (可选，推荐，需安装系统 tesseract)
- easyocr (可选)
- rapidfuzz (可选)
- opencv-python
- numpy

安装示例：
pip install pillow pytesseract easyocr rapidfuzz opencv-python numpy

注意：请根据你的系统修改 TESSERACT_CMD 路径（如果使用 pytesseract）。
"""

import os
import threading
import csv
import time
import tkinter as tk
from tkinter import ttk, filedialog, messagebox
from tkinter.scrolledtext import ScrolledText
from PIL import Image

# 额外依赖
import cv2
import numpy as np

# 尝试导入 pytesseract / easyocr / rapidfuzz
USE_PYTESSERACT = False
USE_EASYOCR = False
USE_RAPIDFUZZ = False

# 请按需修改为你的 tesseract 可执行路径（正斜杠或原始字符串均可）
TESSERACT_CMD = 'D:/OCR/tesseract.exe'

try:
    import pytesseract
    USE_PYTESSERACT = True
except Exception:
    USE_PYTESSERACT = False

try:
    import easyocr
    USE_EASYOCR = True
except Exception:
    USE_EASYOCR = False

try:
    from rapidfuzz import fuzz
    USE_RAPIDFUZZ = True
except Exception:
    USE_RAPIDFUZZ = False

import difflib


# ---------- 图像预处理与文本区域增强函数 ----------

def pil_to_cv(img_pil):
    return cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)


def cv_to_pil(img_cv):
    return Image.fromarray(cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB))


def upscale_pil(img_pil, scale=2):
    w, h = img_pil.size
    return img_pil.resize((int(w * scale), int(h * scale)), Image.LANCZOS)


def preprocess_button_for_white_text(cv_img):
    """针对白字在彩色按钮上的处理：反色 -> 高斯模糊 -> 自适应二值化 -> 开操作"""
    gray = cv2.cvtColor(cv_img, cv2.COLOR_BGR2GRAY)
    inv = 255 - gray
    blur = cv2.GaussianBlur(inv, (3, 3), 0)
    th = cv2.adaptiveThreshold(blur, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                               cv2.THRESH_BINARY, 11, 2)
    kernel = np.ones((2, 2), np.uint8)
    opened = cv2.morphologyEx(th, cv2.MORPH_OPEN, kernel)
    return opened


def preprocess_small_gray_text(cv_img):
    """针对浅灰细字的处理：CLAHE -> 锐化 -> 自适应二值化"""
    gray = cv2.cvtColor(cv_img, cv2.COLOR_BGR2GRAY)
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
    eq = clahe.apply(gray)
    gaussian = cv2.GaussianBlur(eq, (0, 0), 2)
    unsharp = cv2.addWeighted(eq, 1.5, gaussian, -0.5, 0)
    th = cv2.adaptiveThreshold(unsharp, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                               cv2.THRESH_BINARY, 15, 8)
    return th


def detect_red_button_region(cv_img):
    """尝试用 HSV 颜色阈值检测红色按钮的轮廓并返回最大轮廓的 bbox（带边距）。
    如果找不到则返回 None。"""
    hsv = cv2.cvtColor(cv_img, cv2.COLOR_BGR2HSV)
    # 红色在 HSV 上有两个区间
    lower1 = np.array([0, 60, 40])
    upper1 = np.array([10, 255, 255])
    lower2 = np.array([160, 60, 40])
    upper2 = np.array([179, 255, 255])
    mask1 = cv2.inRange(hsv, lower1, upper1)
    mask2 = cv2.inRange(hsv, lower2, upper2)
    mask = cv2.bitwise_or(mask1, mask2)
    # 清理噪声
    kernel = np.ones((3, 3), np.uint8)
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=2)
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return None
    # 找最大轮廓
    c = max(contours, key=cv2.contourArea)
    x, y, w, h = cv2.boundingRect(c)
    # 忽略太小的结果
    h_img, w_img = cv_img.shape[:2]
    if w < 20 or h < 10 or w < 0.05 * w_img:
        return None
    pad_x = int(w * 0.1)
    pad_y = int(h * 0.2)
    x0 = max(0, x - pad_x)
    y0 = max(0, y - pad_y)
    x1 = min(w_img, x + w + pad_x)
    y1 = min(h_img, y + h + pad_y)
    return (x0, y0, x1, y1)


def detect_button_candidates(cv_img):
    """基于边缘和轮廓检测潜在的按钮区域（不依赖颜色），返回若干 bbox 列表"""
    gray = cv2.cvtColor(cv_img, cv2.COLOR_BGR2GRAY)
    # 边缘检测
    edges = cv2.Canny(gray, 50, 150)
    # 膨胀以连接边缘
    kernel = np.ones((3, 3), np.uint8)
    edges = cv2.dilate(edges, kernel, iterations=1)
    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    h_img, w_img = cv_img.shape[:2]
    candidates = []
    for c in contours:
        x, y, w, h = cv2.boundingRect(c)
        area = w * h
        if area < 500:  # 忽略太小
            continue
        # 长宽比过滤：按钮通常较宽
        ratio = w / (h + 1e-6)
        if ratio < 1.5:
            continue
        # 忽略几乎全图的轮廓
        if w > 0.9 * w_img and h > 0.6 * h_img:
            continue
        # 验证该区域是否较“有色彩”（避免纯黑白噪声）
        hsv = cv2.cvtColor(cv_img[y:y+h, x:x+w], cv2.COLOR_BGR2HSV)
        sat = hsv[:, :, 1]
        mean_sat = float(np.mean(sat))
        if mean_sat < 20:  # 太灰或者低饱和度可能不是按钮
            continue
        pad_x = int(w * 0.05)
        pad_y = int(h * 0.15)
        x0 = max(0, x - pad_x)
        y0 = max(0, y - pad_y)
        x1 = min(w_img, x + w + pad_x)
        y1 = min(h_img, y + h + pad_y)
        candidates.append((x0, y0, x1, y1))
    # 按面积从大到小排序
    candidates.sort(key=lambda b: (b[2]-b[0])*(b[3]-b[1]), reverse=True)
    return candidates


def ocr_try_multiple(cv_region, languages=None):
    """对给定的 cv 图像区域尝试多种预处理与 OCR 参数，返回识别结果的集合（非空去重）。"""
    results = set()
    pil_region = cv_to_pil(cv_region)
    scales = [2, 3]
    configs = ['--psm 7 --oem 3', '--psm 8 --oem 3', '--psm 6 --oem 3']
    for s in scales:
        up = upscale_pil(pil_region, scale=s)
        up_cv = pil_to_cv(up)
        # 变换集合
        variants = []
        variants.append(up_cv)
        variants.append(preprocess_button_for_white_text(up_cv))
        # 反转然后阈值
        gray = cv2.cvtColor(up_cv, cv2.COLOR_BGR2GRAY)
        inv = 255 - gray
        _, th_inv = cv2.threshold(inv, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        variants.append(th_inv)
        # 自适应阈值
        th_adaptive = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                            cv2.THRESH_BINARY, 11, 2)
        variants.append(th_adaptive)
        # CLAHE
        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
        eq = clahe.apply(gray)
        variants.append(eq)

        for var in variants:
            for cfg in configs:
                try:
                    if USE_PYTESSERACT:
                        pytesseract.pytesseract.tesseract_cmd = TESSERACT_CMD
                        txt = pytesseract.image_to_string(var, config=cfg)
                    elif USE_EASYOCR:
                        reader = easyocr.Reader(languages or ['en'])
                        # easyocr 接受 PIL 或 numpy
                        txt = "/n".join(reader.readtext(cv_to_pil(var), detail=0))
                    else:
                        txt = ''
                except Exception:
                    txt = ''
                txt_norm = normalize_text(txt)
                if txt_norm:
                    results.add(txt.strip())
    return results


# ---------- OCR 主函数（包含区域增强） ----------

def ocr_image(path, languages=None, try_enhance=True):
    """对图片做 OCR，返回识别出的纯文本字符串。
    如果 try_enhance=True，会尝试检测按钮/底部提示等并做局部增强识别，最后把所有识别到的文本合并返回。
    """
    texts = []

    # 先原始全图 OCR
    if USE_PYTESSERACT:
        try:
            pytesseract.pytesseract.tesseract_cmd = TESSERACT_CMD
        except Exception:
            pass

    try:
        if USE_PYTESSERACT:
            lang = '+'.join(languages) if languages else None
            try:
                whole = pytesseract.image_to_string(Image.open(path), lang=lang)
            except Exception:
                whole = ''
        elif USE_EASYOCR:
            reader = easyocr.Reader(languages or ['en'])
            whole = "/n".join(reader.readtext(path, detail=0))
        else:
            whole = ''
    except Exception:
        whole = ''

    if whole:
        texts.append(('full', whole))

    if not try_enhance:
        return '/n'.join(t for _, t in texts)

    # 读取为 cv 图用于局部检测与增强
    pil = Image.open(path).convert('RGB')
    cv_img = pil_to_cv(pil)
    h_img, w_img = cv_img.shape[:2]

    # 1) 尝试检测红色按钮并对按钮区域做针对性增强
    btn_bbox = detect_red_button_region(cv_img)
    if btn_bbox is not None:
        x0, y0, x1, y1 = btn_bbox
        btn_region = cv_img[y0:y1, x0:x1]
        btn_texts = ocr_try_multiple(btn_region, languages=languages)
        for t in btn_texts:
            texts.append(('button', t))

    # 2) 基于形状的候选按钮检测（捕捉非标准颜色的按钮，如渐变或粉色）
    candidates = detect_button_candidates(cv_img)
    # 仅尝试前几个候选以节省时间
    for bbox in candidates[:3]:
        x0, y0, x1, y1 = bbox
        region = cv_img[y0:y1, x0:x1]
        cand_texts = ocr_try_multiple(region, languages=languages)
        for t in cand_texts:
            texts.append(('candidate_button', t))

    # 3) 取底部区域作为可能的提示文字区域（比如 modal 底部的灰色提示）
    bottom_h = int(h_img * 0.25)
    bottom_region = cv_img[h_img - bottom_h:h_img, 0:w_img]
    bottom_pil_up = upscale_pil(cv_to_pil(bottom_region), scale=2)
    bottom_cv_up = pil_to_cv(bottom_pil_up)
    pre_bottom = preprocess_small_gray_text(bottom_cv_up)
    try:
        if USE_PYTESSERACT:
            bottom_text = pytesseract.image_to_string(pre_bottom, config='--psm 6 --oem 3')
        elif USE_EASYOCR:
            reader = easyocr.Reader(languages or ['en'])
            bottom_text = "/n".join(reader.readtext(cv_to_pil(pre_bottom), detail=0))
        else:
            bottom_text = ''
    except Exception:
        bottom_text = ''
    if bottom_text:
        texts.append(('bottom', bottom_text))

    # 合并去重：以换行分隔，并尽量保持内容完整
    final_parts = []
    seen = set()
    for _, t in texts:
        tnorm = normalize_text(t)
        if not tnorm:
            continue
        if tnorm in seen:
            continue
        seen.add(tnorm)
        final_parts.append(t)

    return '/n'.join(final_parts)


# ---------- 文字比较工具函数 ----------

def normalize_text(s):
    if s is None:
        return ""
    return ' '.join(s.split()).strip().lower()


def similarity(a, b):
    """返回 0-100 相似度分数"""
    a_n = normalize_text(a)
    b_n = normalize_text(b)
    if not a_n or not b_n:
        return 0
    if USE_RAPIDFUZZ:
        return int(fuzz.ratio(a_n, b_n))
    return int(difflib.SequenceMatcher(None, a_n, b_n).ratio() * 100)


# ---------- GUI 应用（与你现有逻辑一致，仅替换 OCR 实现） ----------

class OCRCompareApp(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title('OCR 文案对比工具')
        self._build_ui()
        self.image_folder = None
        self.csv_folder = None
        self.match_results = []  # 保存导出数据

    def _build_ui(self):
        frm = ttk.Frame(self)
        frm.pack(fill=tk.BOTH, expand=True, padx=8, pady=8)

        btn_row = ttk.Frame(frm)
        btn_row.pack(fill=tk.X)

        ttk.Button(btn_row, text='选择图片文件夹', command=self.select_image_folder).pack(side=tk.LEFT, padx=4)
        ttk.Button(btn_row, text='选择 CSV 文件夹', command=self.select_csv_folder).pack(side=tk.LEFT, padx=4)
        ttk.Button(btn_row, text='开始处理', command=self.start_process).pack(side=tk.LEFT, padx=4)
        ttk.Button(btn_row, text='导出结果为 CSV', command=self.export_results).pack(side=tk.LEFT, padx=4)

        opt_row = ttk.Frame(frm)
        opt_row.pack(fill=tk.X, pady=6)
        ttk.Label(opt_row, text='模糊匹配阈值:').pack(side=tk.LEFT)
        self.threshold_var = tk.IntVar(value=85)
        ttk.Scale(opt_row, from_=50, to=100, variable=self.threshold_var, orient=tk.HORIZONTAL).pack(side=tk.LEFT, fill=tk.X, expand=True, padx=6)
        self.fuzzy_var = tk.BooleanVar(value=True)
        ttk.Checkbutton(opt_row, text='启用模糊匹配(若未完全包含则使用相似度判断)', variable=self.fuzzy_var).pack(side=tk.LEFT, padx=6)

        progress_row = ttk.Frame(frm)
        progress_row.pack(fill=tk.X)
        self.progress = ttk.Progressbar(progress_row, maximum=100)
        self.progress.pack(fill=tk.X, expand=True, side=tk.LEFT, padx=2, pady=2)
        self.status_label = ttk.Label(progress_row, text='就绪')
        self.status_label.pack(side=tk.LEFT, padx=6)

        # 结果表格
        tree_frame = ttk.Frame(frm)
        tree_frame.pack(fill=tk.BOTH, expand=True, pady=6)
        cols = ('image', 'phrase', 'score', 'status', 'detail')
        self.tree = ttk.Treeview(tree_frame, columns=cols, show='headings')
        for c in cols:
            self.tree.heading(c, text=c)
            self.tree.column(c, width=120, anchor=tk.W)
        self.tree.pack(fill=tk.BOTH, expand=True)

        # 日志
        log_frame = ttk.LabelFrame(frm, text='日志')
        log_frame.pack(fill=tk.BOTH, expand=False, pady=6)
        self.logbox = ScrolledText(log_frame, height=10, state=tk.NORMAL)
        self.logbox.pack(fill=tk.BOTH, expand=True)

    def log(self, msg):
        ts = time.strftime('%Y-%m-%d %H:%M:%S')
        line = f'[{ts}] {msg}'
        self.logbox.insert(tk.END, line)
        self.logbox.see(tk.END)
        # 同时写入文件
        with open('ocr_compare.log', 'a', encoding='utf-8') as f:
            f.write(line)

    def select_image_folder(self):
        d = filedialog.askdirectory()
        if d:
            self.image_folder = d
            self.log(f'选择图片文件夹: {d}')

    def select_csv_folder(self):
        d = filedialog.askdirectory()
        if d:
            self.csv_folder = d
            self.log(f'选择 CSV 文件夹: {d}')

    def start_process(self):
        if not self.image_folder or not self.csv_folder:
            messagebox.showwarning('提示', '请先选择图片文件夹和 CSV 文件夹')
            return
        t = threading.Thread(target=self._process_all)
        t.daemon = True
        t.start()

    def _process_all(self):
        images = [f for f in os.listdir(self.image_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff'))]
        total = len(images)
        if total == 0:
            self.log('图片文件夹中没有发现图片')
            return
        self.match_results = []
        for idx, imgname in enumerate(images, start=1):
            base, _ = os.path.splitext(imgname)
            csvpath = os.path.join(self.csv_folder, base + '.csv')
            imgpath = os.path.join(self.image_folder, imgname)
            if not os.path.exists(csvpath):
                self.log(f'没有找到对应 CSV: {csvpath} (跳过)')
                continue
            self.log(f'处理 {imgname} 对应 {os.path.basename(csvpath)}')
            ocr_text = ocr_image(imgpath)
            if not ocr_text:
                self.log(f'警告：图片 {imgname} OCR 结果为空')

            # 读取 csv 文案，一行一条 (兼容多列，合并每行)
            phrases = []
            try:
                with open(csvpath, newline='', encoding='utf-8') as cf:
                    reader = csv.reader(cf)
                    for row in reader:
                        # 合并每行的列作为一句话
                        if not row:
                            continue
                        text = ' '.join(row).strip()
                        if text:
                            phrases.append(text)
            except UnicodeDecodeError:
                # 尝试 GBK
                with open(csvpath, newline='', encoding='gbk', errors='ignore') as cf:
                    reader = csv.reader(cf)
                    for row in reader:
                        if not row:
                            continue
                        text = ' '.join(row).strip()
                        if text:
                            phrases.append(text)

            # 对每条文案进行检查
            for ph in phrases:
                ph_norm = normalize_text(ph)
                ocr_norm = normalize_text(ocr_text)
                passed = False
                score = 0
                detail = ''
                if ph_norm and ph_norm in ocr_norm:
                    passed = True
                    score = 100
                    detail = '完全匹配(包含)'
                else:
                    if self.fuzzy_var.get():
                        score = similarity(ph, ocr_text)
                        if score >= self.threshold_var.get():
                            passed = True
                            detail = f'模糊匹配通过 (阈值 {self.threshold_var.get()})'
                        else:
                            detail = f'模糊匹配不通过 (得分 {score})'
                    else:
                        detail = '未通过（未开启模糊匹配）'

                status = '通过' if passed else '未通过'
                self.tree.insert('', tk.END, values=(imgname, ph, score, status, detail))
                self.match_results.append({'image': imgname, 'phrase': ph, 'score': score, 'status': status, 'detail': detail})
                self.log(f'[{imgname}] 文案: "{ph}" -> {status} (score={score}) 详情: {detail}')

            pct = int(idx / total * 100)
            self.progress['value'] = pct
            self.status_label.config(text=f'进度: {idx}/{total} ({pct}%)')
            self.update_idletasks()

        self.progress['value'] = 100
        self.status_label.config(text='完成')
        self.log('全部处理完成')

    def export_results(self):
        if not self.match_results:
            messagebox.showinfo('提示', '当前没有结果可以导出')
            return
        f = filedialog.asksaveasfilename(defaultextension='.csv', filetypes=[('CSV 文件', '*.csv')])
        if not f:
            return
        with open(f, 'w', newline='', encoding='utf-8') as of:
            writer = csv.DictWriter(of, fieldnames=['image', 'phrase', 'score', 'status', 'detail'])
            writer.writeheader()
            for r in self.match_results:
                writer.writerow(r)
        self.log(f'导出结果到 {f}')
        messagebox.showinfo('导出完成', f'导出到 {f}')


if __name__ == '__main__':
    app = OCRCompareApp()
    app.geometry('1000x700')
    app.mainloop()
